import numpy as np
import scipy.stats as sps
from numpy.linalg import norm as l2_norm
import matplotlib.pyplot as plt
from tqdm.notebook import tqdm
from IPython.display import clear_output
import plotly.graph_objects as go

def gauss_seidel_solver(A, b, max_iter=100_000, plot_error=False, mode=None):
    x = np.zeros_like(b)
    errs_hist = []
    err = 1
    eps = 1e-3
    
    for it in range(max_iter):     
        if mode == 'symmetric':
            if it % 2:
                for i in range(len(x)):
                    s1 = A[i, 0:i] @ x[0:i, :]
                    s2 = A[i, i+1:] @ x[i+1: , :]

                    x[i] = (b[i] - s1 - s2) / A[i, i]
            else:
                for i in range(len(x) - 1, -1, -1):
                    s1 = A[i, 0:i] @ x[0:i, :]
                    s2 = A[i, i+1:] @ x[i+1:, :]

                    x[i] = (b[i] - s1 - s2) / A[i, i]
        elif mode == 'gauss':
            x_old = x.copy()
            for i in range(len(x)):
                s1 = A[i, 0:i] @ x_old[0:i, :]
                s2 = A[i, i+1:] @ x_old[i+1: , :]

                x[i] = (b[i] - s1 - s2) / A[i, i]  
        else:
            for i in range(len(x)):
                s1 = A[i, 0:i] @ x[0:i, :]
                s2 = A[i, i+1:] @ x[i+1: , :]

                x[i] = (b[i] - s1 - s2) / A[i, i]
        
        err = l2_norm(A @ x - b)  
        
        errs_hist.append(err)
        
        if err < eps:
            return x
        
        if plot_error and i % 10000 == 0:
            clear_output()
            
            plt.figure(figsize=(4, 4))
            plt.plot(errs_hist)
            plt.show()
        
    print('Method failed to converge')
    
    return x

gauss_errs = []
seidel_errs = []
seidel_symm_errs = []

for size in tqdm(range(2, 201)):
    A = np.random.randn(size, size)
    A += np.diag(A.sum(axis=0) + A.sum(axis=1) + 100)
    
    b = np.random.randn(A.shape[1], 1)
    
    gauss_errs.append(l2_norm(A @ gauss_seidel_solver(A, b, mode='gauss') - b))
    seidel_errs.append(l2_norm(A @ gauss_seidel_solver(A, b, mode='seidel') - b))
    seidel_symm_errs.append(l2_norm(A @ gauss_seidel_solver(A, b, mode='symmetric') - b))

  0%|          | 0/199 [00:00<?, ?it/s]

fig = go.Figure()

fig.add_trace(go.Scatter(y=gauss_errs, name='Gauss'))
fig.add_trace(go.Scatter(y=seidel_errs, name='Seidel'))
fig.add_trace(go.Scatter(y=seidel_symm_errs, name='Symmetric Seidel'))

fig.update_layout(title='Diagonally dominant matrix',
                   xaxis_title='Matrix size',
                   yaxis_title='L2-norm of error')

fig.show()

gauss_errs_p = []
seidel_errs_p = []
seidel_symm_errs_p = []

for size in tqdm(range(2, 201)):
    A = np.diag(np.random.randint(size // 2, size, size))  
    U = sps.ortho_group.rvs(size)
    A = U @ A @ U.T
    
    b = np.random.randn(A.shape[1], 1)
    
    gauss_errs_p.append(l2_norm(A @ gauss_seidel_solver(A, b, mode='gauss') - b))
    seidel_errs_p.append(l2_norm(A @ gauss_seidel_solver(A, b, mode='seidel') - b))
    seidel_symm_errs_p.append(l2_norm(A @ gauss_seidel_solver(A, b, mode='symmetric') - b))

  0%|          | 0/199 [00:00<?, ?it/s]

import plotly.graph_objects as go

fig = go.Figure()

fig.add_trace(go.Scatter(y=gauss_errs_p, name='Gauss'))
fig.add_trace(go.Scatter(y=seidel_errs_p, name='Seidel'))
fig.add_trace(go.Scatter(y=seidel_symm_errs_p, name='Symmetric Seidel'))

fig.update_layout(title='Symmetric positive definite matrix',
                  xaxis_title='Matrix size',
                  yaxis_title='L2-norm of error')

fig.show()